#!/usr/bin/env python
import sys
import time
import py_trees


class ContextSwitch(py_trees.behaviour.Behaviour):
    """ """

    def __init__(self, name: str = "ContextSwitch"):
        """Initialise with a behaviour name."""
        super(ContextSwitch, self).__init__(name)
        self.feedback_message = "no context"

    def initialise(self) -> None:
        """Backup and set a new context."""
        self.logger.debug("%s.initialise()[switch context]" % (self.__class__.__name__))
        # Some actions that:
        #   1. retrieve the current context from somewhere
        #   2. cache the context internally
        #   3. apply a new context
        self.feedback_message = "new context"

    def update(self) -> py_trees.common.Status:
        """Just returns RUNNING while it waits for other activities to finish."""
        self.logger.debug(
            "%s.update()[RUNNING][%s]"
            % (self.__class__.__name__, self.feedback_message)
        )
        return py_trees.common.Status.RUNNING

    def terminate(self, new_status: py_trees.common.Status) -> None:
        """Restore the context with the previously backed up context."""
        self.logger.debug(
            "%s.terminate()[%s->%s][restore context]"
            % (self.__class__.__name__, self.status, new_status)
        )
        # Some actions that:
        #   1. restore the cached context
        self.feedback_message = "restored context"


def create_root() -> py_trees.behaviour.Behaviour:
    """
    Create the root behaviour and it's subtree.
    Returns:
        the root behaviour
    """
    root = py_trees.composites.Parallel(
        name="Parallel", policy=py_trees.common.ParallelPolicy.SuccessOnOne()
    )
    context_switch = ContextSwitch(name="Context")
    sequence = py_trees.composites.Sequence(name="Sequence", memory=True)
    for job in ["Action 1", "Action 2"]:
        success_after_two = py_trees.behaviours.StatusQueue(
            name=job,
            queue=[py_trees.common.Status.RUNNING, py_trees.common.Status.RUNNING],
            eventually=py_trees.common.Status.SUCCESS,
        )
        sequence.add_child(success_after_two)
    root.add_child(context_switch)
    root.add_child(sequence)
    return root


def main() -> None:
    py_trees.logging.level = py_trees.logging.Level.DEBUG
    root = create_root()
    # py_trees.display.render_dot_tree(root)
    root.setup_with_descendants()
    for i in range(1, 6):
        try:
            print("\n--------- Tick {0} ---------\n".format(i))
            root.tick_once()
            print("\n")
            print("{}".format(py_trees.display.unicode_tree(root, show_status=True)))
            time.sleep(1.0)
        except KeyboardInterrupt:
            break
    print("\n")


if __name__ == "__main__":
    main()
